import cv2
import numpy as np
import os
import math

import util


def compress_with_opencv(input_path, output_path, target_kb=20):
    img = cv2.imread(input_path, cv2.IMREAD_UNCHANGED)
    original_size = os.path.getsize(input_path) / 1024

    if original_size <= target_kb:
        cv2.imwrite(output_path, img)
        return

    # 计算初始压缩参数
    quality = 90
    params = [cv2.IMWRITE_PNG_COMPRESSION, quality]

    while quality >= 0 and original_size > target_kb:
        # 尝试保存
        cv2.imwrite(output_path, img, params)
        current_size = os.path.getsize(output_path) / 1024

        if current_size <= target_kb:
            break

        # 调整质量
        quality -= 10
        params = [cv2.IMWRITE_PNG_COMPRESSION, max(0, quality)]

    # 如果仍大于目标，调整尺寸
    if os.path.getsize(output_path) / 1024 > target_kb:
        scale_factor = math.sqrt(target_kb / (os.path.getsize(output_path) / 1024))
        new_width = int(img.shape[1] * scale_factor)
        new_height = int(img.shape[0] * scale_factor)
        img = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_AREA)
        cv2.imwrite(output_path, img)

    final_size = os.path.getsize(output_path) / 1024
    print(f"压缩完成，最终大小: {final_size:.2f}KB")


# 使用示例

for i in range(0, 13):
    print(i)
    compress_with_opencv(fr"D:\wx_game\images\pkq\frame_000{util.get_day(i)}.png", fr"D:\wx_game\images\pkq\ff{i}.png", target_kb=40)

# compress_with_opencv(r"D:\wx_game\images\one\frame_00000.png", r"D:\wx_game\images\one\f.png", target_kb=40)
